import sys  
sys.path.append(r'code')
from feature_eng.alpha_101.alpha101 import Alphas101
import os 
import pandas as pd
from tqdm import tqdm
from joblib import Parallel, delayed
import time 

class Calc_101():
    
    def __init__(self, columns_root=None, save_root=None, price_root=None, value_root=None) -> None:
        self.price_root = r'data/stock_data/daily/price' if price_root is None else  price_root
        self.value_root = r'data/stock_data/daily/market_value' if value_root is None else  value_root
        self.columns_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip' if columns_root is None else columns_root        
        self.save_root = r'data/cmodty/my_feature/feature_data_101' if save_root is None else save_root     
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.column = pd.read_pickle(self.columns_root).columns
            
    def calc_alpha(self, func, name='alpha_ser_0'):
        try :
            t1 = time.time()
            lst = []
            for code in self.column:
                price_file = os.path.join(self.price_root,f'{code}.pkl.gzip')
                # value_file = os.path.join(self.value_root,f'{code}.pkl.gzip')
                # price = pd.read_pickle(price_file)
                # price = price[price.index>'2005-01-01']
                # value = pd.read_pickle(value_file)
                # value = value[['total_mv']]
                # value = value[value.index>'2005-01-01']
                # df_i = pd.concat([price,value], axis=1)
                df_i = pd.read_pickle(price_file)
                alpha_i = pd.DataFrame(index = df_i.index) 
                ary = func(name, df_i)
                alpha_i[code] = ary
                lst.append(alpha_i)
            df = pd.concat(lst, axis=1)
            df.columns = self.column
            df = df[df.index>'2005-01-01']
            save_file = os.path.join(self.save_root, f'{name}.pkl.gzip')
            df.to_pickle(save_file)
            print(f"{name} calculating cost time:{(time.time()-t1):.2f}s")
        except:
            with open(r"code\feature_eng\alpha_series\info.txt", 'a', encoding="utf-8") as f:
                    f.write(f"\n{name}")
            print(f"{name} calculating faild")

    def eval_alpha_101(self, name, df_i):
        alp1 = Alphas101(df_i)
        ary = eval(f'alp1.{name}')()
        return ary

def calc_main_101():
    C = Calc_101()
    alph_lst = list(range(1,102))  #缺失 'alpha_179',
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_101, f'alpha{i}') for i in alph_lst)


from feature_eng.alpha_101.alpha_indneu import Alphas_ind, Alphas_ind9
class Ind_main():
    
    def __init__(self, save_root=None) -> None:
        self.save_root = r'data/cmodty/my_feature/feature_data_ind' if save_root is None else save_root     
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)

    def calc_alpha(self, func, name='alpha_ser_0'):
        try:
            t1 = time.time()
            df = func(name)
            df = df[df.index>'2005-01-01']
            save_file = os.path.join(self.save_root, f'{name}.pkl.gzip')
            df.to_pickle(save_file)
            print(f"{name} calculating cost time:{(time.time()-t1):.2f}s")
        except:
            print(f"{name} calculating failed")


    def eval_alpha_ind(self, name):
        alp1 = Alphas_ind()
        ary = eval(f'alp1.{name}')()
        return ary

def calc_main_alpha_ind():
    # #0: 17
    C = Ind_main()
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_ind, f'alpha_n{i}') for i in range(122))

    # for i in range(1,18):
    #     C.calc_alpha(C.eval_alpha_ind, f'alpha_n{i}')
  
class Ind9_main():
    
    def __init__(self, save_root=None) -> None:
        self.save_root = r'data/cmodty/my_feature/feature_data_ind9' if save_root is None else save_root     
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)

    def calc_alpha(self, func, name='alpha_ser_0'):
        try:
            t1 = time.time()
            df = func(name)
            df = df[df.index>'2005-01-01']
            save_file = os.path.join(self.save_root, f'{name}.pkl.gzip')
            df.to_pickle(save_file)
            print(f"{name} calculating cost time:{(time.time()-t1):.2f}s")
        except:
            print(f"{name} calculating failed")


    def eval_alpha_ind9(self, name):
        alp1 = Alphas_ind9()
        ary = eval(f'alp1.{name}')()
        return ary

def calc_main_alpha_ind9():
    # #0: 17
    C = Ind9_main()
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_ind9, f'alpha_9n{i}') for i in range(1,192))

    # for i in range(1,192):
    #     C.calc_alpha(C.eval_alpha_ind9, f'alpha_9n{i}')
    # C.calc_alpha(C.eval_alpha_ind9, 'alpha_9n181')   
# calc_main_101()  
# calc_main_alpha_ind() 
# calc_main_alpha_ind9()    